#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
第二步：股票筛选系统 - 优化版本
专门为每周2-3%收益目标设计的高效筛选器
"""

import sys
import os
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import warnings
import time
warnings.filterwarnings('ignore')

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from real_data_fetcher import RealDataFetcher
from config import SCREENING_CRITERIA

class OptimizedStockScreener:
    """优化版股票筛选器"""
    
    def __init__(self):
        self.data_fetcher = RealDataFetcher()
        self.criteria = SCREENING_CRITERIA
        
        # 筛选条件
        self.screening_config = {
            'market_cap_min': 40,      # 最小市值40亿（硬性条件）
            'market_cap_max': 3000,    # 最大市值3000亿
            'price_min': 4.0,          # 最低价格4元（硬性条件）
            'price_max': 25.0,         # 最高价格25元（硬性条件）
            'pe_min': 0.1,             # 最小PE>0（硬性条件，确保盈利）
            'pe_max': 80,              # 最大PE<80（避免过高估值）
            'min_volume': 1000000,     # 最小成交量100万
            'target_stocks': 15        # 目标筛选股票数量
        }
        
        # 预设的优质股票候选池（基于市值和流动性）
        self.candidate_pool = [
            # 银行股（稳定分红）
            {'code': '000001', 'name': '平安银行', 'sector': '银行'},
            {'code': '000002', 'name': '万科A', 'sector': '地产'},
            {'code': '600036', 'name': '招商银行', 'sector': '银行'},
            {'code': '002142', 'name': '宁波银行', 'sector': '银行'},
            {'code': '601318', 'name': '中国平安', 'sector': '保险'},
            
            # 科技股（成长性好）
            {'code': '000063', 'name': '中兴通讯', 'sector': '通信'},
            {'code': '002415', 'name': '海康威视', 'sector': '安防'},
            {'code': '002230', 'name': '科大讯飞', 'sector': '软件'},
            {'code': '300059', 'name': '东方财富', 'sector': '金融服务'},
            
            # 消费股（防御性强）
            {'code': '000858', 'name': '五粮液', 'sector': '白酒'},
            {'code': '600519', 'name': '贵州茅台', 'sector': '白酒'},
            {'code': '002304', 'name': '洋河股份', 'sector': '白酒'},
            
            # 制造业（周期性）
            {'code': '600309', 'name': '万华化学', 'sector': '化工'},
            {'code': '002594', 'name': '比亚迪', 'sector': '汽车'},
            {'code': '300750', 'name': '宁德时代', 'sector': '电池'},
            
            # 公用事业（稳定收益）
            {'code': '600900', 'name': '长江电力', 'sector': '电力'},
            {'code': '000725', 'name': '京东方A', 'sector': '面板'},
            
            # 其他优质股票
            {'code': '300015', 'name': '爱尔眼科', 'sector': '医疗'},
            {'code': '000776', 'name': '广发证券', 'sector': '券商'},
            {'code': '600196', 'name': '复星医药', 'sector': '医药'},
        ]
        
        # 统计数据
        self.screening_stats = {
            'total_candidates': 0,
            'passed_basic_filter': 0,
            'final_selected': 0,
            'success_rate': 0.0,
            'start_time': None,
            'end_time': None,
            'processing_time': 0.0
        }
        
    def print_header(self, title):
        """打印标题"""
        print(f"\n{'='*80}")
        print(f"📊 {title}")
        print(f"{'='*80}")
        
    def print_criteria(self):
        """打印筛选条件"""
        print(f"📋 筛选条件（硬性条件）:")
        print(f"   • 市值: >{self.screening_config['market_cap_min']}亿元")
        print(f"   • 股价: {self.screening_config['price_min']}-{self.screening_config['price_max']}元")
        print(f"   • 市盈率: >{self.screening_config['pe_min']}（确保盈利）")
        print(f"   • 成交量: >{self.screening_config['min_volume']:,}手")
        print(f"   • 目标数量: {self.screening_config['target_stocks']}只")
        print(f"   • 候选股池: {len(self.candidate_pool)}只")
        
    def get_stock_info_batch(self, stock_codes):
        """批量获取股票信息"""
        results = []
        
        for i, stock_code in enumerate(stock_codes, 1):
            print(f"   [{i:2d}/{len(stock_codes)}] 获取 {stock_code} 信息...")
            
            try:
                # 获取基本信息
                info = self.data_fetcher.get_stock_basic_info(stock_code)
                if info:
                    results.append(info)
                else:
                    print(f"      ❌ 获取失败")
                    
            except Exception as e:
                print(f"      ❌ 异常: {e}")
                continue
                
            # 适当延迟以避免过于频繁的请求
            time.sleep(0.5)
            
        return results
        
    def apply_hard_criteria(self, stock_info):
        """应用硬性筛选条件"""
        if not stock_info:
            return False, "信息缺失"
            
        # 检查市值（硬性条件）
        market_cap = stock_info.get('market_cap', 0)
        if market_cap <= self.screening_config['market_cap_min']:
            return False, f"市值过小({market_cap:.1f}亿<{self.screening_config['market_cap_min']}亿)"
            
        # 检查股价（硬性条件）
        price = stock_info.get('current_price', 0)
        if price < self.screening_config['price_min'] or price > self.screening_config['price_max']:
            return False, f"股价不符合({price:.2f}元，要求{self.screening_config['price_min']}-{self.screening_config['price_max']}元)"
            
        # 检查PE（硬性条件）
        pe_ratio = stock_info.get('pe_ratio', 0)
        if pe_ratio <= self.screening_config['pe_min']:
            return False, f"PE过低或亏损({pe_ratio:.1f}，要求>{self.screening_config['pe_min']})"
            
        if pe_ratio > self.screening_config['pe_max']:
            return False, f"PE过高({pe_ratio:.1f}，要求<{self.screening_config['pe_max']})"
            
        # 检查成交量
        volume = stock_info.get('volume', 0)
        if volume < self.screening_config['min_volume']:
            return False, f"成交量不足({volume:,}，要求>{self.screening_config['min_volume']:,})"
            
        return True, "符合条件"
        
    def calculate_screening_score(self, stock_info, candidate_info):
        """计算筛选评分（针对周收益优化）"""
        score = 0
        
        # 基础分（20分）
        score += 20
        
        # 1. 市值评分（25分）- 偏好中等市值
        market_cap = stock_info.get('market_cap', 0)
        if 100 <= market_cap <= 800:
            score += 25  # 中等市值，流动性好
        elif 80 <= market_cap <= 100 or 800 <= market_cap <= 1500:
            score += 20  # 中小市值或大市值
        elif 40 <= market_cap <= 80:
            score += 15  # 小市值
        else:
            score += 10  # 超大市值
            
        # 2. 价格评分（20分）- 偏好中等价格
        price = stock_info.get('current_price', 0)
        if 8 <= price <= 20:
            score += 20  # 中等价格，波动适中
        elif 6 <= price <= 8 or 20 <= price <= 25:
            score += 15  # 偏低或偏高价格
        elif 4 <= price <= 6:
            score += 10  # 低价股
        else:
            score += 5   # 超低价或超高价
            
        # 3. 估值评分（20分）- 偏好合理估值
        pe_ratio = stock_info.get('pe_ratio', 0)
        if 8 <= pe_ratio <= 25:
            score += 20  # 合理估值
        elif 5 <= pe_ratio <= 8 or 25 <= pe_ratio <= 40:
            score += 15  # 偏低或偏高估值
        elif 1 <= pe_ratio <= 5:
            score += 10  # 低估值
        else:
            score += 5   # 高估值
            
        # 4. 行业评分（10分）- 偏好稳定行业
        sector = candidate_info.get('sector', '')
        if sector in ['银行', '保险', '电力', '白酒']:
            score += 10  # 防御性行业
        elif sector in ['通信', '医疗', '医药']:
            score += 8   # 成长性行业
        elif sector in ['软件', '电池', '汽车']:
            score += 6   # 科技行业
        else:
            score += 4   # 其他行业
            
        # 5. 流动性评分（5分）
        volume = stock_info.get('volume', 0)
        if volume >= 10000000:  # 1000万以上
            score += 5
        elif volume >= 5000000:  # 500万以上
            score += 4
        elif volume >= 1000000:  # 100万以上
            score += 3
        else:
            score += 1
            
        return min(score, 100)  # 最高100分
        
    def screen_stocks(self):
        """执行股票筛选"""
        self.print_header("第二步：股票筛选系统 - 优化版本")
        
        # 记录开始时间
        self.screening_stats['start_time'] = datetime.now()
        
        # 打印筛选条件
        self.print_criteria()
        
        # 准备候选股票代码列表
        candidate_codes = [stock['code'] for stock in self.candidate_pool]
        self.screening_stats['total_candidates'] = len(candidate_codes)
        
        print(f"\n🔍 开始筛选 {len(candidate_codes)} 只候选股票...")
        
        # 批量获取股票信息
        print(f"\n📊 获取股票基本信息...")
        stock_infos = self.get_stock_info_batch(candidate_codes)
        
        print(f"✅ 成功获取 {len(stock_infos)}/{len(candidate_codes)} 只股票信息")
        
        # 应用硬性筛选条件
        print(f"\n🔧 应用硬性筛选条件...")
        passed_stocks = []
        
        for stock_info in stock_infos:
            stock_code = stock_info.get('code', '')
            
            # 找到对应的候选股票信息
            candidate_info = next((c for c in self.candidate_pool if c['code'] == stock_code), {})
            
            # 应用硬性条件
            passed, reason = self.apply_hard_criteria(stock_info)
            
            if passed:
                # 计算评分
                score = self.calculate_screening_score(stock_info, candidate_info)
                
                # 合并信息
                final_info = {
                    **stock_info,
                    **candidate_info,
                    'screening_score': score,
                    'passed_reason': reason
                }
                
                passed_stocks.append(final_info)
                
                print(f"   ✅ {stock_code} {stock_info.get('name', '')} - 评分: {score:.1f}")
                print(f"      价格: {stock_info.get('current_price', 0):.2f}元, "
                      f"市值: {stock_info.get('market_cap', 0):.1f}亿, "
                      f"PE: {stock_info.get('pe_ratio', 0):.1f}")
                
            else:
                print(f"   ❌ {stock_code} {stock_info.get('name', '')} - {reason}")
                
        self.screening_stats['passed_basic_filter'] = len(passed_stocks)
        
        # 按评分排序
        passed_stocks.sort(key=lambda x: x['screening_score'], reverse=True)
        
        # 选择前N只股票
        final_stocks = passed_stocks[:self.screening_config['target_stocks']]
        self.screening_stats['final_selected'] = len(final_stocks)
        
        # 记录结束时间
        self.screening_stats['end_time'] = datetime.now()
        self.screening_stats['processing_time'] = (
            self.screening_stats['end_time'] - self.screening_stats['start_time']
        ).total_seconds()
        
        # 生成筛选报告
        self.generate_screening_report(final_stocks)
        
        return final_stocks
        
    def generate_screening_report(self, selected_stocks):
        """生成筛选报告"""
        print(f"\n📊 筛选结果报告")
        print("=" * 80)
        
        stats = self.screening_stats
        
        # 统计信息
        print(f"📈 筛选统计:")
        print(f"   • 候选股票数: {stats['total_candidates']}")
        print(f"   • 通过硬性条件: {stats['passed_basic_filter']}")
        print(f"   • 最终选中: {stats['final_selected']}")
        print(f"   • 筛选成功率: {stats['passed_basic_filter']/stats['total_candidates']*100:.1f}%")
        print(f"   • 处理时间: {stats['processing_time']:.1f}秒")
        
        # 筛选结果
        if len(selected_stocks) > 0:
            print(f"\n🏆 筛选结果 TOP{len(selected_stocks)}:")
            print(f"{'排名':<4} {'代码':<8} {'名称':<12} {'行业':<8} {'价格':<8} {'市值':<8} {'PE':<6} {'评分':<6}")
            print("-" * 80)
            
            for i, stock in enumerate(selected_stocks, 1):
                print(f"{i:<4} {stock['code']:<8} {stock['name']:<12} "
                      f"{stock.get('sector', ''):<8} {stock['current_price']:<8.2f} "
                      f"{stock['market_cap']:<8.1f} {stock['pe_ratio']:<6.1f} {stock['screening_score']:<6.1f}")
                      
            # 分析结果
            print(f"\n🎯 筛选结果分析:")
            
            # 价格分布
            prices = [s['current_price'] for s in selected_stocks]
            print(f"   • 价格区间: {min(prices):.2f}-{max(prices):.2f}元")
            
            # 市值分布
            market_caps = [s['market_cap'] for s in selected_stocks]
            print(f"   • 市值区间: {min(market_caps):.1f}-{max(market_caps):.1f}亿")
            
            # PE分布
            pe_ratios = [s['pe_ratio'] for s in selected_stocks]
            print(f"   • PE区间: {min(pe_ratios):.1f}-{max(pe_ratios):.1f}")
            
            # 行业分布
            sectors = {}
            for stock in selected_stocks:
                sector = stock.get('sector', '其他')
                sectors[sector] = sectors.get(sector, 0) + 1
            print(f"   • 行业分布: {sectors}")
            
            # 推荐前5名
            print(f"\n🎯 本周推荐买入（前5名）:")
            for i, stock in enumerate(selected_stocks[:5], 1):
                print(f"   {i}. {stock['code']} {stock['name']} "
                      f"({stock.get('sector', '')}) - 评分: {stock['screening_score']:.1f}")
                
        else:
            print(f"\n⚠️ 没有股票通过筛选条件")
            print(f"💡 建议:")
            print(f"   • 检查网络连接和数据源")
            print(f"   • 适当放宽筛选条件")
            print(f"   • 增加候选股票池")
            
    def save_screening_results(self, selected_stocks, filename="screening_results.csv"):
        """保存筛选结果到CSV文件"""
        if selected_stocks:
            # 准备数据
            data = []
            for stock in selected_stocks:
                data.append({
                    '排名': selected_stocks.index(stock) + 1,
                    '股票代码': stock['code'],
                    '股票名称': stock['name'],
                    '行业': stock.get('sector', ''),
                    '当前价格': stock['current_price'],
                    '市值(亿)': stock['market_cap'],
                    '市盈率': stock['pe_ratio'],
                    '筛选评分': stock['screening_score'],
                    '筛选时间': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                })
            
            # 保存到CSV
            df = pd.DataFrame(data)
            df.to_csv(filename, index=False, encoding='utf-8-sig')
            print(f"\n💾 筛选结果已保存到 {filename}")
            
            return filename
        else:
            print(f"\n⚠️ 没有筛选结果需要保存")
            return None
            
    def get_screening_stats(self):
        """获取筛选统计信息"""
        return self.screening_stats.copy()

if __name__ == "__main__":
    # 创建优化版筛选器
    screener = OptimizedStockScreener()
    
    # 执行筛选
    results = screener.screen_stocks()
    
    # 保存结果
    if results:
        screener.save_screening_results(results)
    
    # 获取统计信息
    stats = screener.get_screening_stats()
    
    print(f"\n🎉 第二步：股票筛选系统优化完成！")
    print(f"   • 筛选出 {len(results)} 只符合条件的股票")
    print(f"   • 处理时间: {stats['processing_time']:.1f}秒")
    print(f"   • 准备进入第三步：技术分析系统")
